from rknn.api import RKNN

def convert_onnx_to_rknn():
    # 创建RKNN对象
    rknn = RKNN()
    
    # 配置模型参数，针对RK3576进行优化
    print('--> Config model')
    rknn.config(
        target_platform="rk3576",
        # 修正op_target配置，使用正确的格式
        optimization_level=3
    )
    print('done')
    
    # 加载ONNX模型
    print('--> Loading ONNX model')
    ret = rknn.load_onnx(
        model='yolov8n.onnx'
    )
    if ret != 0:
        print('Load ONNX model failed!')
        exit(ret)
    print('done')
    
    # 构建模型
    print('--> Building model')
    ret = rknn.build(do_quantization=False)
    if ret != 0:
        print('Build model failed!')
        exit(ret)
    print('done')
    
    # 导出RKNN模型
    print('--> Export RKNN model')
    ret = rknn.export_rknn('yolov8n_rk3576.rknn')
    if ret != 0:
        print('Export RKNN model failed!')
        exit(ret)
    print('done')
    
    # 释放资源
    rknn.release()

if __name__ == '__main__':
    convert_onnx_to_rknn()
    print('Model conversion completed successfully!')
    